# websocket_manager.py
from fastapi import WebSocket
from typing import List, Dict, Optional
import json
import asyncio
import redis.asyncio as redis


class ConnectionManager:
    """
    WebSocket 连接管理器（支持跨后端通信）
    
    功能：
    1. 管理本地 WebSocket 连接
    2. 通过 Redis Pub/Sub 实现跨后端消息传递
    3. 支持点对点通信和广播
    """
    
    def __init__(self, redis_url: str = "redis://localhost:6379"):
        # 本地 WebSocket 连接管理
        self.active_connections: List[WebSocket] = []
        self.client_connections: Dict[str, WebSocket] = {}
        
        # Redis 配置
        self.redis_url = redis_url
        self.redis_client: Optional[redis.Redis] = None
        self.pubsub = None
        self.channel_name = "customer_service_messages"  # 客服消息频道
        
        # 后台任务
        self._listener_task = None
        self._is_listening = False

    async def init_redis(self):
        """初始化 Redis 连接"""
        if self.redis_client is None:
            try:
                self.redis_client = await redis.from_url(
                    self.redis_url, 
                    encoding="utf-8", 
                    decode_responses=True
                )
                self.pubsub = self.redis_client.pubsub()
                await self.pubsub.subscribe(self.channel_name)
                print(f"✅ Redis 已连接，订阅频道: {self.channel_name}")
                
                # 启动消息监听（放在后台任务中，不阻塞）
                if not self._is_listening:
                    self._listener_task = asyncio.create_task(self._redis_listener())
                    print(f"🎧 Redis 监听任务已创建")
                    # 给监听器一点时间启动，但不等待
                    await asyncio.sleep(0.1)
                    
            except Exception as e:
                print(f"⚠️  Redis 连接失败: {str(e)}")
                print("⚠️  将使用单机模式（无法跨后端通信）")

    async def _redis_listener(self):
        """监听 Redis 消息并转发到本地 WebSocket"""
        self._is_listening = True
        print("🎧 [管理端] Redis 监听器已启动")
        try:
            async for message in self.pubsub.listen():
                if message["type"] == "message":
                    try:
                        data = json.loads(message["data"])
                        target_client = data.get("target_client")
                        msg_content = data.get("message")
                        msg_type = data.get("type", "direct")  # direct / broadcast / user_message
                        
                        print(f"📨 [管理端] 从 Redis 收到消息: type={msg_type}, target={target_client}")
                        
                        if msg_type == "user_message":
                            # 来自业务端的用户消息 - 广播给所有管理端客服
                            await self._local_broadcast(msg_content)
                            print(f"📢 [管理端] 已将用户消息广播给所有管理端客服")
                            
                        elif msg_type == "broadcast":
                            # 广播消息
                            await self._local_broadcast(msg_content)
                            
                        elif target_client:
                            # 点对点消息
                            await self._local_send(target_client, msg_content)
                            
                    except json.JSONDecodeError as e:
                        print(f"⚠️  无法解析 Redis 消息: {message['data']}, error={e}")
                    except Exception as e:
                        print(f"⚠️  处理 Redis 消息时出错: {e}")
        except asyncio.CancelledError:
            print("🛑 [管理端] Redis 监听任务已停止")
        except Exception as e:
            print(f"❌ [管理端] Redis 监听器异常: {e}")
            import traceback
            traceback.print_exc()
        finally:
            self._is_listening = False
            print("🎧 [管理端] Redis 监听器已退出")

    async def connect(self, websocket: WebSocket, client_id: str):
        """接受新连接"""
        print(f"🔷 正在接受 WebSocket 连接: {client_id}")
        try:
            await websocket.accept()
            print(f"🔷 WebSocket.accept() 成功: {client_id}")
        except Exception as e:
            print(f"❌ WebSocket.accept() 失败: {client_id}, error={e}")
            raise
            
        self.active_connections.append(websocket)
        self.client_connections[client_id] = websocket
        print(f"✅ 客户端 {client_id} 已连接，当前连接数：{len(self.active_connections)}")

    def disconnect(self, websocket: WebSocket, client_id: str):
        """移除断开连接"""
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
        if client_id in self.client_connections:
            self.client_connections.pop(client_id)
        print(f"🔌 客户端 {client_id} 已断开，当前连接数：{len(self.active_connections)}")

    async def _local_broadcast(self, message: str):
        """本地广播（不通过 Redis）"""
        # ✅ 确保 message 是字符串
        if isinstance(message, dict):
            message_str = json.dumps(message, ensure_ascii=False)
        else:
            message_str = str(message)
        
        disconnected = []
        for client_id, connection in self.client_connections.items():
            try:
                await connection.send_text(message_str)
            except Exception as e:
                print(f"⚠️  广播消息失败 ({client_id}): {str(e)}")
                disconnected.append((connection, client_id))
        
        # 清理断开的连接
        for ws, cid in disconnected:
            self.disconnect(ws, cid)

    async def _local_send(self, client_id: str, message: str):
        """本地发送（不通过 Redis）"""
        # ✅ 确保 message 是字符串
        if isinstance(message, dict):
            message_str = json.dumps(message, ensure_ascii=False)
        else:
            message_str = str(message)
        
        if client_id in self.client_connections:
            try:
                await self.client_connections[client_id].send_text(message_str)
                print(f"✅ 本地发送消息给 {client_id}")
            except Exception as e:
                print(f"⚠️  本地发送失败 ({client_id}): {str(e)}")
                ws = self.client_connections.get(client_id)
                if ws:
                    self.disconnect(ws, client_id)

    async def broadcast(self, message: str):
        """广播消息（通过 Redis 跨后端）"""
        if isinstance(message, dict):
            message_str = json.dumps(message, ensure_ascii=False)
        else:
            message_str = str(message)
        
        # 本地广播
        await self._local_broadcast(message_str)
        
        # 通过 Redis 广播到其他后端
        if self.redis_client:
            try:
                redis_msg = json.dumps({
                    "type": "broadcast",
                    "message": message_str
                }, ensure_ascii=False)
                await self.redis_client.publish(self.channel_name, redis_msg)
            except Exception as e:
                print(f"⚠️  Redis 广播失败: {str(e)}")

    async def send_to_client(self, client_id: str, message, use_redis: bool = True):
        """
        发消息给指定客户端（支持跨后端）
        
        Args:
            client_id: 客户端ID
            message: 消息内容（字符串或字典）
            use_redis: 是否使用 Redis 跨后端发送（默认 True）
        """
        # 转换消息格式
        if isinstance(message, dict):
            message_str = json.dumps(message, ensure_ascii=False)
        else:
            message_str = str(message)
        
        # 1. 先尝试本地发送
        if client_id in self.client_connections:
            await self._local_send(client_id, message_str)
            return
        
        # 2. 如果本地没有该连接，通过 Redis 发送到其他后端
        if use_redis and self.redis_client:
            try:
                redis_msg = json.dumps({
                    "type": "direct",
                    "target_client": client_id,
                    "message": message_str
                }, ensure_ascii=False)
                await self.redis_client.publish(self.channel_name, redis_msg)
                print(f"📤 通过 Redis 发送消息给 {client_id}")
            except Exception as e:
                print(f"⚠️  Redis 发送失败: {str(e)}")
        else:
            print(f"⚠️  客户端 {client_id} 未连接，当前在线: {list(self.client_connections.keys())}")

    async def publish_to_redis(self, message_type: str, target_client: str = None, content: dict = None):
        """
        直接发布消息到 Redis（用于跨后端通信）
        
        Args:
            message_type: 消息类型 "admin_to_user" 或其他
            target_client: 目标客户端ID
            content: 消息内容（字典）
        """
        if self.redis_client:
            try:
                # 构造 Redis 消息
                # content 是字典，直接使用，不要再次序列化
                redis_msg = json.dumps({
                    "type": message_type,
                    "target_client": target_client,
                    "message": content  # ✅ 直接使用，不再序列化
                }, ensure_ascii=False)
                await self.redis_client.publish(self.channel_name, redis_msg)
                print(f"📤 [管理端] 发布消息到 Redis: type={message_type}, target={target_client}")
            except Exception as e:
                print(f"⚠️  Redis 发布失败: {str(e)}")

    async def get_online_clients(self) -> List[str]:
        """获取当前在线的客户端列表"""
        return list(self.client_connections.keys())

    async def is_client_online(self, client_id: str) -> bool:
        """检查客户端是否在线（仅本地检查）"""
        return client_id in self.client_connections

    async def close(self):
        """关闭管理器，清理资源"""
        if self._listener_task and not self._listener_task.done():
            self._listener_task.cancel()
            try:
                await self._listener_task
            except asyncio.CancelledError:
                pass
        
        if self.pubsub:
            await self.pubsub.unsubscribe(self.channel_name)
            await self.pubsub.close()
        
        if self.redis_client:
            await self.redis_client.close()
        
        print("🛑 WebSocket 管理器已关闭")


# 全局实例（需要根据实际情况配置 Redis URL）
wbs = ConnectionManager(redis_url="redis://localhost:6379")